# -*- coding: utf-8 -*-
"""
@Time    : 2024/8/6 14:08 
@Author  : ZhangShenao 
@File    : 1.使用函数回调实现检索器路由.py 
@Desc    : 使用函数回调实现检索器路由

Function Call函数调用是LLM的一项高级功能
它允许用户传递给LLM一个工具或函数列表,并指定函数的名称、功能、参数定义等等
LLM可以根据用户提问,自动识别需要调用的工具或函数,并按照指定的结构化的格式,生成函数调用的参数
通过Function Call,可以保证LLM输出结果的格式化,有效降低Prompt的复杂度
Function Call意味着LLM可以智能选择不同的工具,并且规范化输出
LLM并不会实际调用函数,而是生成函数调用的参数,之后由本地程序根据生成的参数实际发起调用

可以利用Function Call功能来执行相应的规范化输出
可以构建一个假函数,并告知LLM强制调用这个函数,让LLM返回其函数的参数信息

在LangChain中,通过BaseModel.with_structured_output()方法
即可以将BaseModel转换成Function Call,并强制让LLM调用函数

可以利用Function Call功能,实现检索器的路由
定义一个虚假的函数,强制LLM调用这个函数,LLM就会生成该函数的调用参数
程序可以根据生成的结构化的调用参数,进行检索器路由,选择合适的检索器来进行检索

"""
import os
from typing import Literal

import dotenv
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_openai import ChatOpenAI


# 通过Pydantic,定义路由查询结构
class QueryRoute(BaseModel):
    """将用户的查询路由到最相关的数据源"""

    data_source: Literal["python_docs", "java_docs", "go_docs"] = Field(
        description="根据用户提出的问题，选择与问题最相关的数据源"
    )


def route_retriever(route: QueryRoute) -> str:
    """根据路由参数,返回最相关的数据源的检索器名称"""

    data_source = route.data_source.lower()
    if "python_docs" in data_source:
        return "python_doc_retriever"
    elif "java_docs" in data_source:
        return "java_doc_retriever"
    else:
        return "go_doc_retriever"


dotenv.load_dotenv()

# 创建LLM,并通过with_structured_output()方法,指定函数回调
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0,
                 openai_api_base=os.getenv("OPENAI_API_BASE")).with_structured_output(QueryRoute)

# 构造Chain,并调用
prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个擅长将用户问题路由到适当的数据源的专家。\n请根据问题涉及的编程语言，将其路由到相关数据源"),
    ("human", "{question}")
])
chain = prompt | llm | route_retriever
question = """为什么下面的代码不工作了，请帮我检查下：

ApplicationContext ctx = new ClasspathApplicationContext();
bean = ctx.getBean("userService");
"""

# LLM生成的参数为java_docs
# 实际路由到java_doc_retriever
print(chain.invoke({"question": question}))
